/*
**  test_file: hipGraphTest_largeOfKernel.cpp
**  brief:  test code for large of kernel
*/
#include <hip/hip_runtime.h>
#include <chrono>
#include <vector>
#include <unistd.h>
#include <sys/time.h>

#define KNRM "\x1B[0m"
#define KRED "\x1B[31m"
#define KGRN "\x1B[32m"
#define KYEL "\x1B[33m"
#define KBLU "\x1B[34m"
#define KMAG "\x1B[35m"
#define KCYN "\x1B[36m"
#define KWHT "\x1B[37m"

#define failed(...)                                                                                \
    printf("%serror: ", KRED);                                                                     \
    printf(__VA_ARGS__);                                                                           \
    printf("\n");                                                                                  \
    printf("error: TEST FAILED\n%s", KNRM);                                                        \
    abort();

#define HIPCHECK(error)                                                                            \
    {                                                                                              \
        hipError_t localError = error;                                                             \
        if ((localError != hipSuccess) && (localError != hipErrorPeerAccessAlreadyEnabled)) {      \
            printf("%serror: '%s'(%d) from %s at %s:%d%s\n", KRED, hipGetErrorString(localError),  \
                   localError, #error, __FILE__, __LINE__, KNRM);                                  \
            failed("API returned error code.");                                                    \
        }                                                                                          \
    }

#define TESTLOOP                    1
#define TESTSIZE_OF_2_POWER_MIN     8
#define TESTSIZE_OF_2_POWER_MAX     8
#define BLOCK_SIZE_MIN              32
//#define BLOCK_SIZE_MAX              512
#define BLOCK_SIZE_MAX              64
#define GRID_SIZE_MIN               1
#define GRID_SIZE_MAX               64
#define GRAPH_KERNELS_NUM_MIN       6
#define GRAPH_KERNELS_NUM_MAX       15
#define SPACE_KERNEL                0


double getRealTime() {
    struct timeval tp = {0, 0};
    gettimeofday(&tp, NULL);
    return (double) (tp.tv_sec * 1000000 + tp.tv_usec);
}

__global__ void add(float *d_x, float *d_y, float *d_z, size_t N) {

    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
#if (SPACE_KERNEL == 0)
/*				for(int t=0;t<1024;t++)
				{
					int FF = t*t;
				}
*/
        //printf("in kernel_func: add  d_x: %p d_y: %p d_z: %p size: %zu\n", d_x, d_y, d_z, N);
        d_z[i] = d_x[i] + d_y[i];
#endif
    }
}


bool hipGraphNodeKernelLaunch(float *inputVec_h, float *inputVec_2_h,
                              float *inputVec_d, float *inputVec_2_d,
                              float *outputVec_d, 
                              float *result_h,
                              size_t inputSize, 
                              size_t kernelNum, size_t blockSize, size_t gridSize) {
    HIPCHECK(hipMemcpy(inputVec_d, inputVec_h, sizeof(float) * inputSize, hipMemcpyHostToDevice));
    HIPCHECK(hipMemcpy(inputVec_2_d, inputVec_2_h, sizeof(float) * inputSize, hipMemcpyHostToDevice));
    HIPCHECK(hipMemset(outputVec_d, 0, sizeof(float) * inputSize));
    HIPCHECK(hipMemset(result_h, 0, sizeof(float) * inputSize));

    HIPCHECK(hipStreamSynchronize(nullptr));


    hipGraph_t graph;
    hipGraphExec_t graphExec;
    hipStream_t streamForGraph;
    hipGraphNode_t kernelNode;
    hipGraphNode_t lastNode = NULL;
    int depSize = 0;
    hipKernelNodeParams kernelNodeParams = {0};
    memset(&kernelNodeParams, 0, sizeof(kernelNodeParams));

    double start = getRealTime();
    HIPCHECK(hipStreamCreate(&streamForGraph));


    double graphstart = getRealTime();
    HIPCHECK(hipGraphCreate(&graph, 0));

    void *kernelArgs[4] = {(void *)&inputVec_d, (void *)&inputVec_2_d, (void *)&outputVec_d, &inputSize};
    for (int i = 0; i < kernelNum; i++) {
        kernelNodeParams.func = (void *) add;
        kernelNodeParams.gridDim = dim3(gridSize, 1, 1);
        kernelNodeParams.blockDim = dim3(blockSize, 1, 1);
        kernelNodeParams.sharedMemBytes = 0;
        kernelNodeParams.kernelParams = kernelArgs;
        kernelNodeParams.extra = NULL;
        HIPCHECK(hipGraphAddKernelNode(&kernelNode, graph, &lastNode, depSize, &kernelNodeParams));
        lastNode = kernelNode; depSize = 1;
    }

    HIPCHECK(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
    HIPCHECK(hipGraphLaunch(graphExec, streamForGraph));  //execute graph
    double asyncstart = getRealTime();
    HIPCHECK(hipStreamSynchronize(streamForGraph));
    double asyncend = getRealTime();
    HIPCHECK(hipGraphExecDestroy(graphExec));
    HIPCHECK(hipGraphDestroy(graph));
    double graphend = getRealTime();
    HIPCHECK(hipStreamDestroy(streamForGraph));
    double end = getRealTime();

    HIPCHECK(hipMemcpy(result_h, outputVec_d, sizeof(float) * inputSize, hipMemcpyDeviceToHost));

    printf("graph kernel num: %zd, gridSize=%zd, blockSize=%zd, cost=%12lf, graph=%12lf, cost_wa=%12f, graph_wa=%12f \n",
        kernelNum,
        gridSize,
        blockSize,
        end-start,
        graphend - graphstart,
        (end - start) - (asyncend - asyncstart),
        (graphend - graphstart) - (asyncend - asyncstart));


#if (SPACE_KERNEL == 0)
    double result_h_cpu = 0.0;
    for (int i = 0; i < BLOCK_SIZE_MIN; i++) {
        if (result_h[i] != 3.0f) {
            printf("in func: %s result_h[%d] != 3.0f error \n", __func__, i);
            return false;
        }
    }
#endif
    printf("[%s] graph success \n", __func__);
    return true;
}



bool graphTest() {

    float *inputVec_d = NULL;
    float *inputVec_2_d = NULL;
    float *inputVec_h = NULL;
    float *inputVec_2_h = NULL;
    float *outputVec_d = NULL;
    float *result_h = NULL;

    for(unsigned int loopnum = 0; loopnum < TESTLOOP; loopnum++) {

        size_t sizeStart = 0x1u << TESTSIZE_OF_2_POWER_MIN;
        size_t sizeEnd   = 0x1u <<TESTSIZE_OF_2_POWER_MAX;
        size_t size = 0;

        size_t kernelNumStart = 0x1u << GRAPH_KERNELS_NUM_MIN;
        size_t kernelNumEnd = 0x1u << GRAPH_KERNELS_NUM_MAX;
        size_t kernelNum = 0;

        size_t gridSizeStart = GRID_SIZE_MIN;
        size_t gridSizeEnd = GRID_SIZE_MAX;
        size_t gridSize = 0;

        size_t blockSizeStart = BLOCK_SIZE_MIN;
        size_t blockSizeEnd = BLOCK_SIZE_MAX;
        size_t blockSize = 0;

        hipSetDevice(0);

        for (size = sizeStart; size <= sizeEnd; size = size*2 ){ //8,8

            //for (kernelNum = kernelNumStart; kernelNum <= kernelNumEnd; kernelNum += 64) { //	64,16384*2
            //for (kernelNum = kernelNumEnd; kernelNum <= kernelNumEnd; kernelNum += 64) { //	64,16384*2
            //for (kernelNum = kernelNumStart; kernelNum <= kernelNumEnd; kernelNum *= 2) { //	64,16384*2
            //for (kernelNum = 200000; kernelNum <= 200000; kernelNum<1000?kernelNum+=10:kernelNum+=1000) { //	64,16384*2
//            for (kernelNum = 100000000; kernelNum <= 100000000; kernelNum<1000?kernelNum+=10:kernelNum+=1000) { //	64,16384*2
            for (kernelNum = 504; kernelNum <= 1024; kernelNum+=8) { //	64,16384*2
            //for (kernelNum = 10; kernelNum <= 10; kernelNum+=100) { //	64,16384*2
            //for (kernelNum = 5; kernelNum <= 5; kernelNum<1000?kernelNum+=10:kernelNum+=1000) { //	64,16384*2
 
                for (gridSize = gridSizeStart; gridSize <= gridSizeStart; gridSize += 63) { //1,64

                     for (blockSize = blockSizeStart; blockSize <= blockSizeStart; blockSize *= 2) { //64,64
 
			sleep(1);
                        HIPCHECK(hipHostMalloc(&inputVec_h, sizeof(float) * size, hipHostMallocDefault));
                        HIPCHECK(hipHostMalloc(&inputVec_2_h, sizeof(float) * size, hipHostMallocDefault));
                        HIPCHECK(hipHostMalloc(&result_h, sizeof(float) * size, hipHostMallocDefault));

                        for (int i = 0; i < size; i++) {
                            inputVec_h[i] = 1.0f;
                            inputVec_2_h[i] = 2.0f;
                            result_h[i] = 0.0f;
                        }

                        HIPCHECK(hipMalloc(&inputVec_d, sizeof(float) * size));
                        HIPCHECK(hipMalloc(&inputVec_2_d, sizeof(float) * size));
                        HIPCHECK(hipMalloc(&outputVec_d, sizeof(float) * size));

                        bool statusNode = hipGraphNodeKernelLaunch(inputVec_h, inputVec_2_h, inputVec_d, inputVec_2_d, outputVec_d,
                                                                   result_h, size, kernelNum, blockSize, gridSize);
                        HIPCHECK(hipFree(inputVec_d));
                        HIPCHECK(hipFree(inputVec_2_d));
                        HIPCHECK(hipFree(outputVec_d));
                        HIPCHECK(hipHostFree(inputVec_h));
                        HIPCHECK(hipHostFree(inputVec_2_h));
                        HIPCHECK(hipHostFree(result_h));
                        hipDeviceSynchronize();

                        if (!statusNode) {
                            printf("Failed during hip graph Kernels \n");
                            exit(1);
                        }
                        inputVec_d = NULL;
                        inputVec_2_d = NULL;
                        inputVec_h = NULL;
                        inputVec_2_h = NULL;
                        outputVec_d = NULL;
                        result_h = NULL;

                        //hipDeviceSynchronize();
                    }

                }
                
            }

        }
    }
    return true;
}

int main(int argc, char **argv) {
    graphTest();
    return 0;
}
